import mmcv
import numpy as np
from mmseg.registry import TRANSFORMS
from mmseg.datasets.transforms.loading import LoadAnnotations


@TRANSFORMS.register_module()
class LoadMultiVesselAnnotations(LoadAnnotations):
    """自定义多标签加载组件，用于加载 LCX 和 LAD 两个分支的 mask."""

    def __init__(self,
                 reduce_zero_label: bool = False,
                 imdecode_backend: str = 'pillow'):
        super().__init__(reduce_zero_label=reduce_zero_label,
                         imdecode_backend=imdecode_backend)

    def transform(self, results):
        """加载 LCX 和 LAD 的 mask."""
        lcx_path = results.get('seg_map_path_lcx', None)
        lad_path = results.get('seg_map_path_lad', None)

        if lcx_path is not None:
            results['gt_seg_map_lcx'] = mmcv.imread(
                lcx_path, flag='grayscale', backend=self.imdecode_backend).squeeze()
        if lad_path is not None:
            results['gt_seg_map_lad'] = mmcv.imread(
                lad_path, flag='grayscale', backend=self.imdecode_backend).squeeze()

        # 记录到 seg_fields，方便后续数据增强操作
        results.setdefault('seg_fields', [])
        if 'gt_seg_map_lcx' in results:
            results['seg_fields'].append('gt_seg_map_lcx')
        if 'gt_seg_map_lad' in results:
            results['seg_fields'].append('gt_seg_map_lad')

        return results
